// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pthread.h>                                // pthread_*

#include <cstddef>
#include <memory>
#include <iostream>
#include <tally/impl/agent_group.h>
#include <ktest/ktest.h>

namespace {
    using namespace tally::detail;

    struct Add {
        uint64_t operator()(const uint64_t lhs, const uint64_t rhs) const {
            return lhs + rhs;
        }
    };

    const size_t OPS_PER_THREAD = 2000000;

    class AgentGroupTest : public testing::Test {
    protected:
        typedef std::atomic<uint64_t> agent_type;

        void SetUp() {}

        void TearDown() {}

        static void *thread_counter(void *arg) {
            int id = (int) ((long) arg);
            agent_type *item = AgentGroup<agent_type>::get_or_create_tls_agent(id);
            if (item == NULL) {
                EXPECT_TRUE(false);
                return NULL;
            }
            auto start = turbo::Time::current_time();
            for (size_t i = 0; i < OPS_PER_THREAD; ++i) {
                agent_type *element = AgentGroup<agent_type>::get_or_create_tls_agent(id);
                uint64_t old_value = element->load(std::memory_order_relaxed);
                uint64_t new_value;
                do {
                    new_value = old_value + 2;
                } while (__builtin_expect(!element->compare_exchange_weak(old_value, new_value,
                                                                          std::memory_order_relaxed,
                                                                          std::memory_order_relaxed), 0));
                //element->store(element->load(std::memory_order_relaxed) + 2,
                //               std::memory_order_relaxed);
                //element->fetch_add(2, std::memory_order_relaxed);
            }
            auto end = turbo::Time::current_time();
            auto n_elapsed = turbo::Duration::to_nanoseconds(end - start);
            return (void *) (n_elapsed);
        }
    };

    TEST_F(AgentGroupTest, test_sanity) {
        int id = AgentGroup<agent_type>::create_new_agent();
        ASSERT_TRUE(id >= 0) << id;
        agent_type *element = AgentGroup<agent_type>::get_or_create_tls_agent(id);
        ASSERT_TRUE(element != NULL);
        AgentGroup<agent_type>::destroy_agent(id);
    }

    std::atomic<uint64_t> g_counter(0);

    void *global_add(void *) {
        auto start = turbo::Time::current_time();
        for (size_t i = 0; i < OPS_PER_THREAD; ++i) {
            g_counter.fetch_add(2, std::memory_order_relaxed);
        }
        auto end = turbo::Time::current_time();
        auto n_elapsed = turbo::Duration::to_nanoseconds(end - start);
        return (void *) (n_elapsed);
    }

    TEST_F(AgentGroupTest, test_perf) {
        size_t loops = 100000;
        size_t id_num = 512;
        int ids[id_num];
        for (size_t i = 0; i < id_num; ++i) {
            ids[i] = AgentGroup<agent_type>::create_new_agent();
            ASSERT_TRUE(ids[i] >= 0);
        }
        auto start = turbo::Time::current_time();
        for (size_t i = 0; i < loops; ++i) {
            for (size_t j = 0; j < id_num; ++j) {
                agent_type *agent =
                        AgentGroup<agent_type>::get_or_create_tls_agent(ids[j]);
                ASSERT_TRUE(agent != NULL) << ids[j];
            }
        }
        auto end = turbo::Time::current_time();
        auto n_elapsed = turbo::Duration::to_nanoseconds(end - start);
        KLOG(INFO) << "It takes " << n_elapsed / (loops * id_num)
                  << " ns to get tls agent for " << id_num << " agents";
        for (size_t i = 0; i < id_num; ++i) {
            AgentGroup<agent_type>::destroy_agent(ids[i]);
        }

    }

    TEST_F(AgentGroupTest, test_all_perf) {
        long id = AgentGroup<agent_type>::create_new_agent();
        ASSERT_TRUE(id >= 0) << id;
        pthread_t threads[24];
        for (size_t i = 0; i < TURBO_ARRAYSIZE(threads); ++i) {
            pthread_create(&threads[i], NULL, &thread_counter, (void *) id);
        }
        long totol_time = 0;
        for (size_t i = 0; i < TURBO_ARRAYSIZE(threads); ++i) {
            void *ret;
            pthread_join(threads[i], &ret);
            totol_time += (long) ret;
        }
        KLOG(INFO) << "ThreadAgent takes "
                  << totol_time / (OPS_PER_THREAD * TURBO_ARRAYSIZE(threads));
        totol_time = 0;
        g_counter.store(0, std::memory_order_relaxed);
        for (size_t i = 0; i < TURBO_ARRAYSIZE(threads); ++i) {
            pthread_create(&threads[i], NULL, global_add, (void *) id);
        }
        for (size_t i = 0; i < TURBO_ARRAYSIZE(threads); ++i) {
            void *ret;
            pthread_join(threads[i], &ret);
            totol_time += (long) ret;
        }
        KLOG(INFO) << "Global Atomic takes "
                  << totol_time / (OPS_PER_THREAD * TURBO_ARRAYSIZE(threads));
        AgentGroup<agent_type>::destroy_agent(id);
        //sleep(1000);
    }
} // namespace
